package com.briup.pai.listener;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson2.JSON;
import com.briup.pai.common.constant.ModelConstant;
import com.briup.pai.common.enums.TrainingStatusEnum;
import com.briup.pai.convert.TrainingLabelConvert;
import com.briup.pai.entity.message.ModelTrainResultMessage;
import com.briup.pai.entity.po.Model;
import com.briup.pai.entity.po.Training;
import com.briup.pai.entity.po.TrainingLabel;
import com.briup.pai.service.IModelService;
import com.briup.pai.service.ITrainingLabelService;
import com.briup.pai.service.ITrainingService;
import org.springframework.amqp.rabbit.annotation.Exchange;
import org.springframework.amqp.rabbit.annotation.Queue;
import org.springframework.amqp.rabbit.annotation.QueueBinding;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.aop.framework.AopContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cache.annotation.CacheConfig;
import org.springframework.cache.annotation.CacheEvict;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;

/**
 * 接受python训练完成的消息
 */
@Component
@CacheConfig(cacheNames = ModelConstant.MODEL_CACHE_PREFIX)
public class ModelTrainRabbitListener {

    @Autowired
    private IModelService modelService;
    @Autowired
    private ITrainingLabelService trainingLabelService;
    @Autowired
    private ITrainingService trainingService;
    @Autowired
    private TrainingLabelConvert trainingLabelConvert;

    @RabbitListener(bindings = @QueueBinding(
            value = @Queue(value = "${mq.init.result-queue}",durable = "true"),
            exchange = @Exchange(value = "${mq.init.exchange}"),
            key = "${mq.init.result-routing-key}"
    ))
    @Transactional
    public void receiveInitTrainMessage(String message){
        ModelTrainResultMessage resultMessage = JSON.parseObject(message, ModelTrainResultMessage.class);
        ((ModelTrainRabbitListener)(AopContext.currentProxy())).processTrainingResult(resultMessage);
    }
    @RabbitListener(bindings = @QueueBinding(
            value = @Queue(value = "${mq.optimize.result-queue}",durable = "true"),
            exchange = @Exchange(value = "${mq.optimize.exchange}"),
            key = "${mq.optimize.result-routing-key}"
    ))
    @Transactional
    public void receiveOptimizeTrainMessage(String message){
        ModelTrainResultMessage resultMessage = JSON.parseObject(message, ModelTrainResultMessage.class);
        ((ModelTrainRabbitListener)(AopContext.currentProxy())).processTrainingResult(resultMessage);
    }

    @Transactional
    @CacheEvict(key = "T(com.briup.pai.common.constant.CommonConstant).DETAIL_CACHE_PREFIX+':'+#modelId")
    public void processTrainingResult(ModelTrainResultMessage resultMessage) {
        // 处理训练结果
        // 训练失败，拿到model，进行更新
        Integer modelId = resultMessage.getModelId();
        Model model = modelService.getById(modelId);
        if(ObjectUtil.equal(resultMessage.getIsSuccess(),1)){
            model.setTrainingStatus(TrainingStatusEnum.ERROR.getStatus());
            modelService.updateById(model);
            return;
        }
        // 保存训练的数据
        List<TrainingLabel> trainingLabels =  resultMessage.getLabelResults().stream()
                .map(s -> {
            TrainingLabel trainingLabel = trainingLabelConvert.labelResultMessage2Po(s);
            trainingLabel.setTrainingId(resultMessage.getTrainingId());
            return trainingLabel;
        }).toList();
        trainingLabelService.saveBatch(trainingLabels);

        // 修改训练的状态
        Training training = trainingService.getById(resultMessage.getTrainingId());
        training.setAccuracyRate(resultMessage.getAccuracyRate());
        training.setModelFileAddr(resultMessage.getModelFileAddr());
        trainingService.updateById(training);
        // 修改模型的状态
        model.setLastModelVersion(model.getLastModelVersion()+1);
        model.setTrainingStatus(TrainingStatusEnum.DONE.getStatus());
        model.setAccuracyRate(resultMessage.getAccuracyRate());
        modelService.updateById(model);
    }
}
